#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import logging
import sys
from io import BytesIO

import pandas as pd
from openpyxl import Workbook, load_workbook

from rag.nlp import find_codec


# 定义一个名为 RAGFlowExcelParser 的类，用于处理 Excel 文件解析相关的功能
class RAGFlowExcelParser:

    # 静态方法 _load_excel_to_workbook 接收一个文件类对象作为参数，用于将 Excel 或 CSV 内容加载为工作簿对象
    @staticmethod
    def _load_excel_to_workbook(file_like_object):
        # 如果传入的是 bytes 类型数据，则将其封装成 BytesIO 对象以便后续操作
        if isinstance(file_like_object, bytes):
            file_like_object = BytesIO(file_like_object)

        # Read first 4 bytes to determine file type
        # 将文件指针移到开头，读取前4个字节以判断文件类型（用于区分 Excel 和 CSV）
        file_like_object.seek(0)
        # 读取文件头的前4个字节，保存下来用于后续判断文件类型
        file_head = file_like_object.read(4)
        # 将文件指针恢复到开头位置，方便后续读取整个文件内容
        file_like_object.seek(0)

        # 判断是否是 Excel 文件：检查文件头是否为 ZIP 格式 (PK\x03\x04) 或 OLE 格式 (\xD0\xCF\x11\xE0)
        if not (file_head.startswith(b'PK\x03\x04') or file_head.startswith(b'\xD0\xCF\x11\xE0')):
            # 日志提示不是 Excel 文件，尝试将内容当作 CSV 转换为 Excel 工作簿
            logging.info("****wxy: Not an Excel file, converting CSV to Excel Workbook")

            try:
                # 将文件指针重置到开始位置，准备用 pandas 读取 CSV 数据
                file_like_object.seek(0)
                # 使用 pandas 读取 CSV 数据并转换为 DataFrame
                df = pd.read_csv(file_like_object)
                # 调用私有方法 _dataframe_to_workbook 将 DataFrame 转换为 Excel 工作簿对象
                return RAGFlowExcelParser._dataframe_to_workbook(df)

            except Exception as e_csv:
                # 捕获异常并抛出错误信息，说明 CSV 解析失败
                raise Exception(f"****wxy: Failed to parse CSV and convert to Excel Workbook: {e_csv}")

        # 尝试使用 openpyxl 的 load_workbook 方法加载 Excel 文件内容
        try:
            return load_workbook(file_like_object,data_only= True)
        except Exception as e:
            # 如果 openpyxl 加载失败，记录日志并尝试使用 pandas 读取 Excel
            logging.info(f"****wxy: openpyxl load error: {e}, try pandas instead")
            
            try:
                # 将文件指针重置到开始位置，准备用 pandas 读取 Excel
                file_like_object.seek(0)
                # 使用 pandas 的 read_excel 方法读取 Excel 文件内容为 DataFrame
                df = pd.read_excel(file_like_object)
                # 再次调用 _dataframe_to_workbook 方法将 DataFrame 转换为 Excel 工作簿对象
                return RAGFlowExcelParser._dataframe_to_workbook(df)
            
            except Exception as e_pandas:
                # 如果 pandas 也失败，抛出详细错误信息，包含两个阶段的异常信息
                raise Exception(f"****wxy: pandas.read_excel error: {e_pandas}, original openpyxl error: {e}")

    # 将DataFrame转换为Excel工作簿对象的静态方法
    @staticmethod
    def _dataframe_to_workbook(df):
        # 创建一个新的Excel工作簿实例
        wb = Workbook()
        # 获取活动的工作表并设置标题为"Data"
        ws = wb.active
        ws.title = "Data"

        # 遍历DataFrame的列名，将它们写入工作表的第一行
        for col_num, column_name in enumerate(df.columns, 1):
            ws.cell(row=1, column=col_num, value=column_name)

        # 遍历DataFrame的数据行，从第二行开始写入工作表
        for row_num, row in enumerate(df.values, 2):
            for col_num, value in enumerate(row, 1):
                ws.cell(row=row_num, column=col_num, value=value)

        # 返回包含数据的工作簿对象
        return wb

    # html方法用于解析Excel文件并将其内容以HTML表格形式返回
    def html(self, fnm, chunk_rows=256):
        # 判断输入是否为字节流，否则将其作为文件路径处理
        file_like_object = BytesIO(fnm) if not isinstance(fnm, str) else fnm
        # 调用私有方法加载Excel文件到工作簿对象
        wb = RAGFlowExcelParser._load_excel_to_workbook(file_like_object)
        # 初始化一个空列表用于存储HTML表格块
        tb_chunks = []
        # 遍历工作簿中的每个工作表
        for sheetname in wb.sheetnames:
            # 获取当前工作表
            ws = wb[sheetname]
            # 获取工作表的所有行
            rows = list(ws.rows)
            # 如果工作表为空则跳过
            if not rows:
                continue

            # 构建表格的首行（表头），使用<th>标签包裹
            tb_rows_0 = "<tr>"
            for t in list(rows[0]):
                tb_rows_0 += f"<th>{t.value}</th>"
            tb_rows_0 += "</tr>"

            # 根据指定的每块行数分割表格数据
            for chunk_i in range((len(rows) - 1) // chunk_rows + 1):
                # 初始化HTML表格字符串
                tb = ""
                # 添加表格标题，使用工作表名称作为caption
                tb += f"<table><caption>{sheetname}</caption>"
                # 添加表格首行（表头）
                tb += tb_rows_0
                # 遍历当前块的数据行
                for r in list(
                    rows[1 + chunk_i * chunk_rows: 1 + (chunk_i + 1) * chunk_rows]
                ):
                    # 开始一行表格数据
                    tb += "<tr>"
                    # 遍历当前行的每一列
                    for i, c in enumerate(r):
                        # 如果单元格值为空，则添加空的<td>标签
                        if c.value is None:
                            tb += "<td></td>"
                        else:
                            # 否则添加包含单元格值的<td>标签
                            tb += f"<td>{c.value}</td>"
                    # 结束当前行
                    tb += "</tr>"
                # 结束当前表格块
                tb += "</table>\n"
                # 将当前表格块添加到结果列表中
                tb_chunks.append(tb)

        # 返回所有HTML表格块
        return tb_chunks

    # 调用该实例时，传入文件名或文件内容，解析Excel并返回结构化文本数据
    def __call__(self, fnm):
        # 判断输入是否为字节流，否则将其作为文件路径处理
        file_like_object = BytesIO(fnm) if not isinstance(fnm, str) else fnm
        # 加载Excel文件到工作簿对象
        wb = RAGFlowExcelParser._load_excel_to_workbook(file_like_object)

        # 初始化结果列表，用于存储解析后的每一行文本
        res = []
        # 遍历每个工作表名称
        for sheetname in wb.sheetnames:
            # 获取当前工作表
            ws = wb[sheetname]
            # 获取所有行的列表
            rows = list(ws.rows)
            # 如果当前工作表没有数据则跳过
            if not rows:
                continue
            # 获取表头行，用于字段前缀
            ti = list(rows[0])
            # 遍历从第二行开始的所有数据行
            for r in list(rows[1:]):
                # 初始化字段列表，用于存储当前行解析出的各个字段
                fields = []
                # 遍历当前行的每一个单元格
                for i, c in enumerate(r):
                    # 如果单元格为空，则跳过
                    if not c.value:
                        continue
                    # 获取对应列的表头（若存在），否则为空字符串
                    t = str(ti[i].value) if i < len(ti) else ""
                    # 拼接表头和单元格值，使用中文冒号分隔
                    t += ("：" if t else "") + str(c.value)
                    # 添加拼接好的字段到fields列表中
                    fields.append(t)
                # 使用分号连接所有字段，形成一行文本
                line = "; ".join(fields)
                # 如果工作表名称不包含"sheet"字样，则在行末添加来源工作表名称
                if sheetname.lower().find("sheet") < 0:
                    line += " ——" + sheetname
                # 将构造好的行添加到结果列表中
                res.append(line)
        # 返回最终解析后的文本结果列表
        return res

    # 静态方法：计算给定文件的总行数（支持Excel、CSV、TXT格式）
    @staticmethod
    def row_number(fnm, binary):
        # 如果是Excel文件格式
        if fnm.split(".")[-1].lower().find("xls") >= 0:
            # 加载二进制内容为Excel工作簿
            wb = RAGFlowExcelParser._load_excel_to_workbook(BytesIO(binary))
            # 初始化总行数计数器
            total = 0
            # 遍历每个工作表
            for sheetname in wb.sheetnames:
                # 获取当前工作表
                ws = wb[sheetname]
                # 累加当前工作表的行数
                total += len(list(ws.rows))
            # 返回总行数
            return total

        # 如果是CSV或TXT格式文件
        if fnm.split(".")[-1].lower() in ["csv", "txt"]:
            # 自动检测编码格式
            encoding = find_codec(binary)
            # 解码为文本并忽略非法字符
            txt = binary.decode(encoding, errors="ignore")
            # 返回换行符分割后的行数
            return len(txt.split("\n"))


if __name__ == "__main__":
    psr = RAGFlowExcelParser()
    psr(sys.argv[1])
